// ==================================================================================
// Shared Genomics Project MPI Codebase
// Version 1.0 30/04/2010
//
// (c) 2010 University of Manchester all rights reserved
//
// This file is distributed under the GNU General Public License, Version 2.  
// Please see the file COPYING.txt for more details
// ==================================================================================

/*!
\file 
\brief Implementation of Linear Model (Design Matrix)
*/

#include <math.h>

#include "BitString.h"
#include "linear.h"
#include "pmaths.h"
#include "utils.h"

int linear_model_init(struct linear_model *p) {
	if (p == NULL) return 0; 

	p->miss = NULL;
	p->nMiss = 0;
	p->np = 0;
	p->all_valid = FALSE;

	// Flag variables associated with the design matrix of the logistic model.
	p->mAA = 0;
	p->mAB = 1;
	p->mBB = 2;

	// The n-D data arrays
	p->dimension = 0;
	p->additive = NULL;
	p->haploid = NULL;
	p->xchr = NULL;
	p->nInteraction = 1;

	// 2-D array X (indpendent array)
	p->X = p->Xt = NULL;
	p->nX = 0;
	p->X_max_used = 0;

	// 1-D dependent array.
	p->nind = 0;
	p->Y = p->p = p->V = NULL;
	p->valid = NULL;

	// Initialise the FitLM variables.
	p->coef = p->sig = p->w = p->b = NULL;
	p->S  = p->u = p->v = p->Xt = NULL;
	p->S0 = p->S0_inverse = NULL;
	p->nU = 0;
	p->varY = 0;
	p->chisq = 0;
	p->cluster = FALSE;
	
	return 1;
}

int linear_model_free_attributes(struct linear_model *lm) {
	if (lm == NULL) return 0;

	if (lm->miss != NULL) {
		free(lm->miss);
		lm->miss = NULL;
		lm->nMiss = 0;
	}

	if (lm->valid != NULL) { free(lm->valid); lm->valid = NULL; }

	// Free the n-D results array.
	if (lm->dimension != 0) {
		if (lm->additive != NULL) { free(lm->additive); lm->additive = NULL; }
		if (lm->haploid != NULL) { free(lm->haploid); lm->haploid = NULL;  }
		if (lm->xchr != NULL)  { free(lm->xchr); lm->xchr = NULL; }
	}
	
	// Clean-up the 'X' design matrix component.
	if (lm->X != NULL) {
		free_matrix(&lm->X);
		lm->X = NULL;
		lm->nX = 0;
		lm->X_max_used = 0;
	}

	// Clean-up 'Y' bit of the design matrix
	if (lm->p != NULL) { free(lm->p); lm->p = NULL; }
	if (lm->V != NULL) { free(lm->V); lm->V = NULL; }
	if (lm->Y != NULL) { free(lm->Y); lm->Y = NULL; }
	
	// Clean up the FIT LM variables.
	if (lm->coef != NULL) { free(lm->coef); lm->coef = NULL; }
	if (lm->S != NULL) { free_matrix(&lm->S); lm->S = NULL; }
	if (lm->sig != NULL) { free(lm->sig); lm->sig = NULL; }
	if (lm->w != NULL) { free(lm->w); lm->w = NULL; } 
	if (lm->u != NULL) { free_matrix(&lm->u); lm->u = NULL;  lm->nU = 0; }
	if (lm->v != NULL) { free_matrix(&lm->v); lm->v = NULL; }
	if (lm->Xt != NULL) { free_matrix(&lm->Xt); lm->Xt = NULL; }
	if (lm->b != NULL) { free(lm->b); lm->b = NULL; }
	if (lm->S0 != NULL) { free_matrix(&lm->S0); lm->S0 = NULL; }
	if (lm->S0_inverse != NULL) { free_matrix(&lm->S0_inverse); lm->S0_inverse = NULL; } 
	
	// Clear the remaining index variables.
	lm->np = 0;
	lm->nind = 0; 

	return 1;
}

void linear_model_set_np(struct linear_model *lm, struct selected_options *ops) {
	if (lm == NULL || ops == NULL) return;

	// Number of parameters
	// Epistasis 1 for intercept column, 2 for additive effect considered and 1 for the interaction.
	lm->np = 1 + 2 + 1;

	// Sex effect?
	if (ops->sex_effect) lm->np++;
}

int linear_model_alloc(struct linear_model *lm, int nSamples) {
	int dimension = 2;

	if (lm == NULL) return 0;
	else if (lm->np == 0) return 0;
	else if (lm->miss != NULL) return 0; // Already allocated.
	
	lm->miss = (BOOL*) malloc(sizeof(BOOL) * nSamples);
	if (lm->miss != NULL) lm->nMiss = nSamples;
	lm->additive = (int*) malloc(sizeof(int) * dimension);
	lm->haploid = (BOOL*) malloc(sizeof(BOOL) * dimension);
	lm->xchr = (BOOL*) malloc(sizeof(BOOL) * dimension);
	lm->dimension = dimension; // The n-wise character of the algorithm, at the moment N=2.
	lm->X = create_matrix(nSamples, lm->np, 0);
	if (lm->X != NULL) lm->nX = nSamples;
	lm->Y = create_vector(nSamples, 0);
	lm->p = create_vector(nSamples, 0);
	lm->V = create_vector(nSamples, 0);
	lm->coef = create_vector(lm->np, 0);
	lm->S = create_matrix(lm->np, lm->np, 0);
	lm->sig = create_vector(nSamples, 0);
	lm->w = create_vector(lm->np, 0);
	lm->u = create_matrix(nSamples, lm->np, 0);
	if (lm->u != NULL) lm->nU = nSamples;
	lm->v = create_matrix(lm->np, lm->np, 0);
	lm->Xt = create_matrix(lm->np, nSamples, 0);
	lm->b = create_vector(nSamples, 0);
	lm->S0 = create_matrix(lm->np, lm->np, 0);
	lm->S0_inverse = create_matrix(lm->np, lm->np, 0);
	lm->valid = (BOOL*) malloc(sizeof(BOOL) * lm->np);

	if (lm->miss == NULL || lm->additive == NULL || lm->haploid == NULL ||  lm->xchr == NULL || lm->X == NULL ||
		lm->Y == NULL || lm->p == NULL || lm->V == NULL || 
		lm->coef == NULL || lm->S == NULL || lm->sig == NULL || lm->w == NULL || lm->u == NULL || lm->v == NULL ||
		lm->Xt == NULL || lm->b == NULL || lm->S0 == NULL || lm->S0_inverse == NULL || lm->valid == NULL) 
		return 0;
	
	return 1;
}

// Set phenotype to the 'Y' dependent array.
int linear_model_set_dependent(struct linear_model *lm, struct sample *samples, int nSamples) {
	int i;

	if (lm == NULL || samples == NULL) return 0;
	else if (lm->Y == NULL || nSamples == 0) return 0;

	for (i = 0, lm->nind = 0; i < nSamples; i++)
		if (!lm->miss[i]) lm->Y[lm->nind++] = samples[i].pperson->phenotype;

	return 1;
}

int linear_model_set_variance(struct linear_model *lm) {
	double mean = 0;
	int actualN = 0, i;

	if (lm == NULL) return 0;
	else if (lm->nind == 0 || lm->Y == NULL) return 0;

	lm->varY = 0;

	for (i = 0; i < lm->nind; i++) {
		actualN++;
		mean += lm->Y[i];
	}
	if (actualN == 0) {
		lm->varY = 0;
		return 1;
	}
	mean /= (double) actualN;

	for (i = 0; i < lm->nind; i++) lm->varY += (lm->Y[i] - mean) * (lm->Y[i] - mean);	
	lm->varY /= (double)(actualN - 1);

	if (actualN != lm->nind) return 0;

	return 1;
}

int linear_model_add_additive_snp(struct linear_model *lm, struct chromosomes *chromos, struct small_locus *m, int snp_index, int dimension_index) {
	if (lm == NULL || chromos == NULL || m == NULL || snp_index < 0 || dimension_index < 0)
		return 0;

	lm->additive[dimension_index] = snp_index;

	if (chromos->chr_sex[m[snp_index].chr]) 
		lm->xchr[dimension_index] = TRUE;
	else
		lm->xchr[dimension_index] = FALSE;

	if (chromos->chr_haploid[m[snp_index].chr]) 
		lm->haploid[dimension_index] = TRUE;
	else
		lm->haploid[dimension_index] = FALSE;

	return 1;
}

/*
Build X matrix (including intercept)
Iterate a person at a time, entering only 
valid rows into X (non-missing); also build Y
at the same time.
*/
enum MATHS_RETURN linear_model_build_design_matrix(struct linear_model *lm, char *ipedData, struct sample *samples, int nSamples, int nIpedRecordSizeInBytes, struct selected_options *ops) {
	int i, j;
	MATRIX_T X = NULL;  // Pointer linear model X co-ordinate.
	char *p = NULL;
	int ped_offset = nIpedRecordSizeInBytes * sizeof(char);
	
	if (lm == NULL || ipedData == NULL) return MEMORY_FAILURE;
	
	// Count the number of missing samples in the sample array (if any).
	// to get dimension for the 2-D vector array.
	for (i = 0, X = lm->X, lm->X_max_used = 0, p = ipedData; i < nSamples; i++, p += ped_offset) {
		struct sample *person = &samples[i];
		VECTOR_T trow = NULL; // Current row in the design matrix.
		int c = 1; 
		BOOL skip = FALSE;
		int *s = NULL; // 2-D list of SNP indexes to be considered for pairwise association things.
		int snp;
		
		// Ignore if missing phenotype, or the user set this to missing
		if (lm->miss[i]) continue;
		
		// List of indexes in the design matrix.
		// 0) Intercept      
		// 1) Main effects of SNPs 
		// 2) Interactions of the above

		trow = X[lm->X_max_used];
		trow[0] = 1; // Intercept
		
		// Additive effects (assuming individual-major mode)
		for (j = 0, s = lm->additive, snp = 0; j < lm->dimension; j++, s++, snp++) {
			long bitpos = (*s) * 2;
			unsigned int mask = 0x1; 
			BitString bs = (BitString) p;
			BOOL i1, i2;

			// Get first bit.
			mask = mask << (bitpos % BITBULKSIZE);
			mask = bs[bitpos/BITBULKSIZE] & mask;
			i1 = (mask !=0) ? TRUE : FALSE;
			
			// Get second bit
			bitpos++;
			mask = 0x1; 
			mask = mask << (bitpos % BITBULKSIZE);
			mask = bs[bitpos/BITBULKSIZE] & mask;
			i2 = (mask !=0) ? TRUE : FALSE;

			if (lm->xchr[snp]) {
				// X chromosome coding
				double mA = 0, mB = 1;
				if (ops->xchr_model == 1) {
					mA = 0.5; 
					mB = 1.5;
				} else if ( ops->xchr_model == 2 ) {
					mA = 0;
					mB = 2;
				} else if ( ops->xchr_model > 2 ) {
					mA = 0;
					mB = 1;
				}

				// male
				if ( person->sex ) {
					if (i1) {
						if (!i2) {
							lm->miss[i] = TRUE;
							skip = TRUE;
						}
						else
							trow[c++] = mA; 	      
					} else {
						if (i2) {
							// This should not happen...
							lm->miss[i] = TRUE;
							skip = TRUE;
						}		  
						else
							trow[c++] = mB; 
					}
				} 
				// female x-chromosome
				else {
					if (i1) {
						if (!i2) {
							lm->miss[i] = TRUE;
							skip = TRUE;
						}
						else
							trow[c++] = 0; 	      
					} else {
						if (i2)
							trow[c++] = 1; // het 
						else
							trow[c++] = 2; // hom
					}		  
				}

			}
			else if (lm->haploid[snp]) {
				// Haploid coding
				if (i1) {
					if (!i2) {
						lm->miss[i] = TRUE;
						skip = TRUE;
					} else
						trow[c++] = 0; 	      
				} else {
					if (i2) {
						// This should not happen...
						lm->miss[i] = TRUE;
						skip = TRUE;
					}		  
					else
						trow[c++] = 1; 
				}
			} else {
				// Autosomal coding
				if (i1) {
					if (!i2) {
						lm->miss[i] = TRUE;
						skip = TRUE;
					}
					else
						trow[c++] = lm->mAA;
				}
				else {
					if (i2)
						trow[c++] = lm->mAB; // het 
					else
						trow[c++] = lm->mBB; // hom
				}
			} 
		}
		
		// Flag if need to ignore a row in the 2-D array.
		if (skip) continue; 
		
		// Sex effect (automatically set for X chromosome models)
		// Seems to be associated with GLM analysis only.
		if ( ops->sex_effect ) {
			if (samples[i].sex) trow[c++] = 1;
			else trow[c++] = 0;
		}

		// Interactions 
		for (j = 0; j < lm->nInteraction; j++)
			trow[c++] = trow[lm->interaction[j].p1] * trow[lm->interaction[j].p2];
		
		lm->X_max_used++; // If reach here, all is well so increment flag to use another row in the 'X' 2-D array.
	}
	
	// VIF-based check for multicollinearity
	lm->all_valid = check_vif(lm->X, lm->X_max_used, lm->np);

	// Add Y variable also
	if(!linear_model_set_dependent(lm, samples, nSamples)) return MEMORY_FAILURE; 
	
	return MATHS_SUCCESS;
}

enum MATHS_RETURN linear_model_fitLM(struct linear_model *lm) {
	int i, j;
	const double TOL = 1.0e-13;
	double wmax = 0.0, tmp, thresh, sum;
	VECTOR_T afunc = NULL;
	BOOL callFailed = FALSE;

	if (lm == NULL) return MATHS_FAILURE;
	else if (lm->np == 0 || lm->nind == 0 || !lm->all_valid) return MATHS_FAILURE;
	
	if (!linear_model_set_variance(lm)) return MATHS_FAILURE;

	set_vector(lm->sig, lm->nind, sqrt(1.0/sqrt((double)lm->nind)));
	zero_matrix(lm->u, lm->nind, lm->np);
	zero_matrix(lm->v, lm->np, lm->np);
	
	for (i = 0; i < lm->nind; i++) {
		afunc = lm->X[i];
		tmp = 1.0/lm->sig[i];
		for (j = 0; j < lm->np; j++) lm->u[i][j] = afunc[j] * tmp;     
		lm->b[i] = lm->Y[i] * tmp;
	}
	
	if (!svdcmp(lm->u, lm->nind, lm->np, lm->w, lm->v)) 
		return MATHS_FAILURE;
	
	// Set the threshold values.
	for (j = 0; j < lm->np; j++) if (lm->w[j] > wmax) wmax = lm->w[j];
	thresh = TOL * wmax;
	for (j = 0; j < lm->np; j++) if (lm->w[j] < thresh) lm->w[j] = 0.0;

	// Do a bit of back substiution.
	svbksb(lm->u, lm->nind, lm->np, lm->w, lm->v, lm->b, lm->coef);
	lm->chisq = 0.0;
	for (i = 0; i < lm->nind; i++) {
		afunc = lm->X[i];
		sum = 0.0;
		for (j = 0; j < lm->np; j++) sum += lm->coef[j] * afunc[j];
		lm->chisq += (tmp = (lm->Y[i]-sum)/lm->sig[i], tmp * tmp);
	} 
	
	/*
	Obtain covariance matrix of estimates

	OLS variance estimator = s^2 * ( X'X )^-1
	where s^2 = (1/(N-k)) \sum_i=1^N e_i^2

	Robust cluster variance estimator
	V_cluster = (X'X)^-1 * \sum_{j=1}^{n_C} u_{j}' * u_j * (X'X)^-1 
	where u_j = \sum_j cluster e_i * x_i 

	Above, e_i is the residual for the ith observation and x_i is a
	row vector of predictors including the constant.

	For simplicity, I omitted the multipliers (which are close to 1)
	from the formulas for Vrob and Vclusters.

	The formula for the clustered estimator is simply that of the
	robust (unclustered) estimator with the individual ei*s replaced
	by their sums over each cluster. xi

	Williams, R. L. 2000.  A note on robust variance estimation for
	cluster-correlated data. Biometrics 56: 64


	t ( y - yhat X  ) %*%  ( y - yhat)  / nind - np = variance of residuals 

	// Variance of residuals

	j <- ( t( y- m %*% t(b) ) %*% ( y - m %*% t(b) ) ) / ( N - p ) 
	print( sqrt(kronecker( solve( t(m) %*% m ) , j )  ))

	Calcuate S = (XtX)^-1 
	*/
	
	// get the transpose of 'X'
	for (i = 0;  i < lm->nind; i++)
		for (j = 0; j < lm->np; j++) 
			lm->Xt[j][i] = lm->X[i][j];
	
	zero_matrix(lm->S0, lm->np, lm->np);
	mult_matrix(lm->Xt, lm->np, lm->nind, lm->X, lm->nind, lm->np, lm->S0);
	if (!svd_inverse(lm->S0, lm->S0_inverse, lm->np, lm->np)) 
		return MATHS_FAILURE;
	
	// Standard OLS s^2
	if (!lm->cluster) {
		double sigma= 0.0;

		for (i = 0; i < lm->nind; i++) {
			double partial = 0.0;
			for (j = 0; j < lm->np; j++) partial += lm->coef[j] * lm->X[i][j];
			partial -= lm->Y[i];
			sigma += partial * partial;
		}
		sigma /= lm->nind - lm->np;	

		for (i = 0; i < lm->np; i++)
			for (j = 0; j < lm->np; j++) lm->S[i][j] = lm->S0_inverse[i][j] * sigma;
	}
		
	if (!callFailed) return MATHS_SUCCESS;
	else return MATHS_FAILURE;
}

enum MATHS_RETURN linear_model_validParameters(struct linear_model *lm) {
	int i, j;

	if (lm == NULL) return MATHS_FAILURE;
	// Empty model?
	else if (lm->np == 0 || lm->nind == 0) {
		lm->all_valid = FALSE;
		return MATHS_FAILURE;
	}

	// Check for multicollinearity
	// For each term, see that estimate is not too strongly (r>0.99)
	// correlated with another, starting at last
	for (i = 1; i < lm->np; i++) {
		lm->valid[i] = TRUE;
		if ( lm->S[i][i] < 1e-20 ) 
			lm->valid[i] = lm->all_valid = FALSE; 
		else if ( ! realnum(lm->S[i][i]) ) 
			lm->valid[i] = lm->all_valid = FALSE; 
	}

	if (lm->all_valid) {
		for (i = lm->np - 1; i > 0; i--) {
			for (j = i-1; j >= 0; j--) {	  
				if (lm->S[i][j] / sqrt( lm->S[i][i] * lm->S[j][j] ) > 0.99999) {
					lm->valid[i] = FALSE;
					lm->all_valid = FALSE;
					break;
				}
			}
		}
	} 

	return MATHS_SUCCESS;
}

double linear_model_get_statistic(struct linear_model *lm, int testParameter) {
	if (lm->all_valid) {
		return ( lm->coef[testParameter] * lm->coef[testParameter] ) 
			/ lm->S[testParameter][testParameter];
	}
	else 
		return 0;
}